{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Week 5 Day 4\n",
    "\n",
    "AutoGen Core - Distributed\n",
    "\n",
    "I'm only going to give a Teaser of this!!\n",
    "\n",
    "Partly because I'm unsure how relevant it is to you. If you'd like me to add more content for this, please do let me know.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.messages import TextMessage\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "from autogen_ext.tools.langchain import LangChainToolAdapter\n",
    "from langchain_community.utilities import GoogleSerperAPIWrapper\n",
    "from langchain.agents import Tool\n",
    "from IPython.display import display, Markdown\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "ALL_IN_ONE_WORKER = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start with our Message class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@dataclass\n",
    "class Message:\n",
    "    content: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And now - a host for our distributed runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.runtimes.grpc import GrpcWorkerAgentRuntimeHost\n",
    "\n",
    "host = GrpcWorkerAgentRuntimeHost(address=\"localhost:50051\")\n",
    "host.start() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's reintroduce a tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "serper = GoogleSerperAPIWrapper()\n",
    "langchain_serper =Tool(name=\"internet_search\", func=serper.run, description=\"Useful for when you need to search the internet\")\n",
    "autogen_serper = LangChainToolAdapter(langchain_serper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "instruction1 = \"To help with a decision on whether to use AutoGen in a new AI Agent project, \\\n",
    "please research and briefly respond with reasons in favor of choosing AutoGen; the pros of AutoGen.\"\n",
    "\n",
    "instruction2 = \"To help with a decision on whether to use AutoGen in a new AI Agent project, \\\n",
    "please research and briefly respond with reasons against choosing AutoGen; the cons of Autogen.\"\n",
    "\n",
    "judge = \"You must make a decision on whether to use AutoGen for a project. \\\n",
    "Your research team has come up with the following reasons for and against. \\\n",
    "Based purely on the research from your team, please respond with your decision and brief rationale.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And make some Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Player1Agent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client, tools=[autogen_serper], reflect_on_tool_use=True)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        return Message(content=response.chat_message.content)\n",
    "    \n",
    "class Player2Agent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client, tools=[autogen_serper], reflect_on_tool_use=True)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        return Message(content=response.chat_message.content)\n",
    "    \n",
    "class Judge(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client)\n",
    "        \n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        message1 = Message(content=instruction1)\n",
    "        message2 = Message(content=instruction2)\n",
    "        inner_1 = AgentId(\"player1\", \"default\")\n",
    "        inner_2 = AgentId(\"player2\", \"default\")\n",
    "        response1 = await self.send_message(message1, inner_1)\n",
    "        response2 = await self.send_message(message2, inner_2)\n",
    "        result = f\"## Pros of AutoGen:\\n{response1.content}\\n\\n## Cons of AutoGen:\\n{response2.content}\\n\\n\"\n",
    "        judgement = f\"{judge}\\n{result}Respond with your decision and brief explanation\"\n",
    "        message = TextMessage(content=judgement, source=\"user\")\n",
    "        response = await self._delegate.on_messages([message], ctx.cancellation_token)\n",
    "        return Message(content=result + \"\\n\\n## Decision:\\n\\n\" + response.chat_message.content)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.runtimes.grpc import GrpcWorkerAgentRuntime\n",
    "\n",
    "if ALL_IN_ONE_WORKER:\n",
    "\n",
    "    worker = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker.start()\n",
    "\n",
    "    await Player1Agent.register(worker, \"player1\", lambda: Player1Agent(\"player1\"))\n",
    "    await Player2Agent.register(worker, \"player2\", lambda: Player2Agent(\"player2\"))\n",
    "    await Judge.register(worker, \"judge\", lambda: Judge(\"judge\"))\n",
    "\n",
    "    agent_id = AgentId(\"judge\", \"default\")\n",
    "\n",
    "else:\n",
    "\n",
    "    worker1 = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker1.start()\n",
    "    await Player1Agent.register(worker1, \"player1\", lambda: Player1Agent(\"player1\"))\n",
    "\n",
    "    worker2 = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker2.start()\n",
    "    await Player2Agent.register(worker2, \"player2\", lambda: Player2Agent(\"player2\"))\n",
    "\n",
    "    worker = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker.start()\n",
    "    await Judge.register(worker, \"judge\", lambda: Judge(\"judge\"))\n",
    "    agent_id = AgentId(\"judge\", \"default\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await worker.send_message(Message(content=\"Go!\"), agent_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(response.content))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "await worker.stop()\n",
    "if not ALL_IN_ONE_WORKER:\n",
    "    await worker1.stop()\n",
    "    await worker2.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "await host.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
